package com.csx.governance.traffic.adapter.outbound.webclient;

import com.csx.governance.traffic.adapter.outbound.resttemplate.TrafficRestTemplateInterceptor;
import com.csx.governance.traffic.core.context.TrafficContext;
import com.csx.governance.traffic.core.context.TrafficContextHolder;
import com.csx.governance.traffic.core.context.TrafficInternalContexts;
import com.csx.governance.traffic.core.context.TrafficRequest;
import com.csx.governance.traffic.core.context.webClient.TrafficWebClientRequest;
import com.csx.governance.traffic.core.interceptor.TrafficHandlerInterceptorRunner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.context.support.GenericApplicationContext;

/**
 * @author cuisongxu
 * @date 2023/4/12 22:38
 */
public class TrafficLoadBalancerLifecycle implements LoadBalancerLifecycle<RequestDataContext, ResponseData, ServiceInstance> {

    private final Logger logger = LoggerFactory.getLogger(TrafficLoadBalancerLifecycle.class);

    private final TrafficHandlerInterceptorRunner handlerInterceptorRunner;

    private final GenericApplicationContext applicationContext;

    private ThreadLocal<TrafficWebClientRequest> trafficWebClientRequestThreadLocal = new ThreadLocal<>();

    public TrafficLoadBalancerLifecycle(GenericApplicationContext applicationContext,
                                        TrafficHandlerInterceptorRunner handlerInterceptorRunner) {
        this.applicationContext = applicationContext;
        this.handlerInterceptorRunner = handlerInterceptorRunner;
    }

    @Override
    public boolean supports(Class requestContextClass, Class responseClass, Class serverTypeClass) {
        return LoadBalancerLifecycle.super.supports(requestContextClass, responseClass, serverTypeClass);
    }

    @Override
    public void onStart(Request<RequestDataContext> request) {
        TrafficWebClientRequest trafficWebClientRequest = new TrafficWebClientRequest(request);
        trafficWebClientRequestThreadLocal.set(trafficWebClientRequest);
        handlerInterceptorRunner.preHandle(trafficWebClientRequest, null);
        buildContext(trafficWebClientRequest);
    }

    @Override
    public void onStartRequest(Request request, Response lbResponse) {
        TrafficWebClientRequest trafficWebClientRequest = trafficWebClientRequestThreadLocal.get();
        handlerInterceptorRunner.postHandle(trafficWebClientRequest, null);
    }

    @Override
    public void onComplete(CompletionContext completionContext) {
        TrafficWebClientRequest trafficWebClientRequest = trafficWebClientRequestThreadLocal.get();
        handlerInterceptorRunner.afterCompletion(trafficWebClientRequest, null, completionContext.getThrowable());
    }

    private void buildContext(TrafficRequest request) {
        TrafficContext context = TrafficContextHolder.getContext();
        if(context != null) {
            context.setSubRequest(request);
        } else { // 说明是由内部定时任务发起的外调，而不是由入栈请求导致的外调
            TrafficInternalContexts.set(request);
        }
    }
}
